import torch.nn as nn
import torch.nn.functional as F
import torch

#创建 cnn模型
class Model(nn.Module):
    def __init__(self):
        # 定义父类构造函数，当创建该类的时候会被调用
        super().__init__()
        # 定义了第一个二维【卷积层】(conv1)，它接收1个输入通道(例如灰度图像)
        # 输出16个通道，使用5x5的卷积核，并设置padding=1,即在图像边缘填充一层0，以保证输出大小与输入大小一致(除去了边界部分)
        self.covn1 = nn.Conv2d(1, 16, 5, padding=1)
        # 定义了2x2的最大【池化层】pool,用于下采样，即特征图的尺寸减小一半，同时保留最重要的特征
        self.pool = nn.MaxPool2d((2, 2))    #下采样
        # 定义了第二个二维【卷积层】(conv2),它接收16个输入通道(来自上一个卷积层的输出)
        # 输出32个通道，同样使用5x5的卷积核，并设置padding=1
        self.covn2 = nn.Conv2d(16, 32, 5,padding=1)
        # 定义第一个【全连接层】linear1,这里经过两个卷积层和两次最大池化后，每个样本特征图大小为32*5*5(具体取决于输入尺寸)
        # 这个层将这些特征展平为一维向量，并映射256个特征
        self.linear1 = nn.Linear(32*5*5, 256)   # 16*4*4
        # 定义第二个【全连接层】linear2,它们将256个特征映射到10个类别(十分类问题)
        self.linear2 = nn.Linear(256, 10)

    # --------------#
    #    向前传播
    # --------------#
    def forward(self, input):       #forward 方法定义向前传播的过程，input是模型的输入数据
        # 应用第一个卷积层conv1，然后通过Relu激活函数，这是为了引入非线性，使得模型能够学习更复杂的模式
        x = F.relu(self.covn1(input))
        # 对上一步得到的特征图应用最大池化操作，进一步减少特征图的尺寸
        x = self.pool(x)
        # 应用第二个卷积层conv2，再通过Relu激活函数
        x = F.relu(self.covn2(x))
        # 再次应用最大池化操作，进一步缩小特征图的尺寸
        x = self.pool(x)
        # print(x.size())
        # 使用view方法将特征图展平为一维向量，x.size(0)保持批次维度不变，-1表示自动计算剩余维度的大小，以适应张量的元素数量
        x = x.view(x.size(0), -1)   #或者 x = x.view(-1, 32*5*5)
        # 将展平后的特征向量通过第一个全连接层linear1，并再次应用Relu激活函数
        x = F.relu(self.linear1(x))     # nn.CroEntropyloss()损失函数需要的是未激活的输出
        # 返回向前传播的最终结果，即对每个样本预测的10个类别的得分
        x = self.linear2(x)
        return x
# #打印模型结构
# model = Model()
# print(model)

#------------------------------------------------------------------------#
#   关于16*4*4如何计算的问题，如果你输入的图像大小是28*28(例如MNIST数据集)，那么
#   - 经过第一次卷积(5x5卷积核，padding=1)，输出大小仍然是28x28